import sys
from pathlib import Path

base_path = str(Path(__file__).resolve().parent.parent)
sys.path.append(base_path)
from pydantic import BaseModel, Field
from langchain_core.output_parsers import JsonOutputParser
from langchain_core.prompts import PromptTemplate
from create_llm import create_llm


class Person(BaseModel):
    name: str = Field(description="用户姓名")
    age: int = Field(description="用户年龄")
    hobbies: list[str] = Field(description="用户爱好列表")


json_parser = JsonOutputParser(pydantic_object=Person)
# instructions = json_parser.get_format_instructions()
# print(instructions)

prompt = PromptTemplate.from_template(
    "根据以下信息和格式要求提取用户信息：\n{info}\n{format_instructions}"
)

llm = create_llm()

chain = prompt | llm | json_parser
temp = chain.invoke(
    {
        "info": "一个喜欢足球和音乐的30岁程序员",
        "format_instructions": json_parser.get_format_instructions(),
    }
)
print(temp)
